"""
#===============================================================================
# License
#===============================================================================
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

#===============================================================================
Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Repository Path:    
$$HeadURL$$

Created on:           7/27/2009
Last Revised on:      $$Date$$
Version:              $$Revision$$
ID:                   $$Id$$
Author:               $$Author$$ :: (bryan)

#===============================================================================


#===============================================================================
# Documentation
#===============================================================================

"""

#===============================================================================
# Imports
#===============================================================================
import JT
from JT.Optimizer import RegularStepGradientDescentOptimizer
import numpy as np
import my_auth
import time
import csv
import os
import sys
from socket import gethostname
from datetime import datetime
import time


class Test(object):
    """
    this module will make it easy to construct registration tests. 
    """

    def __init__(self):    
        # First setup the JointTrack data structure
        self.data = JT.Data.gold_standard()    
#===============================================================================
#    Setup DRR
#===============================================================================      
        self.drr = JT.DRR()
        self.drr._renWin.SetOffScreenRendering(0)
        
#===============================================================================
# Setup Registration
#===============================================================================
        self.reg=JT.MultiIterationRegistration()
        # Set additional arguments for cost function
        self.reg._plot = True
        self.reg._debug = True
        
    def run(self,n_pose=50,method='bpc',vert=0,proj=0,body=0):   
#===============================================================================
# Setup data 
#===============================================================================
        self.drr.SetExternalCalibration(self.data.Frames[vert].FixedImages[proj].calibration)
        self.drr.RemoveAllVolumes()
        self.drr.AddVolume(self.data.Frames[vert].Volumes[0].data)
        self.reg.SetFixedImage(self.data.Frames[vert].FixedImages[proj].data.GetImage(0))
        self.reg.SetMovingImageGenerator(self.drr)
        self.reg.SetRegionOfInterest(self.data.Frames[vert].FixedImages[proj].ROIs[body].data)
        # set Metric
        self.reg.SetImageMetric(method)
        center = self.data.Frames[vert].Volumes[0].BodyCenters[body]
        self.drr.AddBodyCenter(center)
        self.drr.AddSphereAtBodyCenter(center)
        
        # Translate Volume
#        for dim in range(3):
#            for i in range(-10,11,2):
#                pose = np.zeros(6)
#                pose[dim] = i
#                value = self.reg.GetValue(pose)
        # Rotate Volume
        for dim in range(3,6):
            for i in np.arange(-np.pi/18,np.pi/18,np.pi/90.0):
                pose = np.zeros(6)
                pose[dim] = i
                value = self.reg.GetValue(pose)   

if __name__ == "__main__":
    
    test = Test()
    #method_list = ['nmi', 'mi', 'ncc', 'ecc', 'mse', 'ncc', 'gd', 'gc']
    #method_list = ['bpc']
    method_list = ['bpc','nmi','mi','ncc','ecc','ncc','gd','gc','mse']
                   
    vert = [0,1]
    proj = [0,1]
    body = [[0,1,2],[0,1,2,3,4]]
    #body=[[1],[1]]
    if 1:
        v=1
        p=1
        b=4
        method = 'mi'
        tag = 'Method: %s, Vert: %s, Proj: %s, Body: %s' %\
              (method,str(v),str(p),str(b))
        print tag
        test.run(method=method,vert=v,proj=p,body=b)
        #test.reg.ShowFinalResults()
    else:
        for method in method_list:
            for v in vert:
                for p in proj:
                    for b in body[v]:
                        for n_pose in range(200):               
                            # to label the resulting image that will be saved.
                            tag = 'Method: %s, Vert: %s, Proj: %s, Body: %s, Pose: %i' %\
                              (method,str(v),str(p),str(b),n_pose)
                            print tag
                            test.run(n_pose=n_pose,method=method,
                                     vert=v,proj=p,body=b)
                        
                        
                        